import torch
import torch.utils
import torch.utils.data
from typing import Any

class MyTensorDataset(torch.utils.data.Dataset):
    
    def __init__(self, *args) -> None:
        super().__init__()
        self.tensors = (*args, )
        xlen = len(self)
        self.tensors = tuple(tensor[:xlen] for tensor in self.tensors)
        
    def __getitem__(self, index) -> Any:
        return tuple([tensor[index] for tensor in self.tensors])
    
    def __len__(self) -> int:
        return min(len(tensor) for tensor in self.tensors)


if '__main__' == __name__:
    
    import random
    from PyCmpltrtok.common import sep
    
    random.seed(66)
    
    seq1 = torch.Tensor(list(range(1, 20)))
    seq10 = seq1[:15] * 10
    seq100 = seq1[:17] * 100
    ds = MyTensorDataset(seq1, seq10, seq100)
    
    xlen = len(ds)
    print(f'len = {xlen}')
    
    sep('Random access')
    idxs = random.choices(range(xlen), k=10)
    idxs = [idx - 5 for idx in idxs]
    print(idxs)
    for idx in idxs:
        print(idx, ds[idx])
        
    sep('Iteration')
    for idx in range(-5, xlen):
        print(idx, ds[idx])        
        
    for n_worker in range(1, 5 + 1):
        sep(f'num_workers={n_worker}')
        print(list(torch.utils.data.DataLoader(ds, num_workers=n_worker, worker_init_fn=None)))
        